import torch
from model import MattingNetwork
from inference import convert_video


def main():
    # 设置参数
    file_name = "1.mp4"
    variant = 'resnet50'
    # variant = 'mobilenetv3'
    checkpoint = r"C:\www\RobustVideoMatting\checkpoint\rvm_"+variant+".pth"
    device = 'cuda'
    # device = 'cpu'
    input_source = "C:\www\RobustVideoMatting\input\\" + file_name
    output_type = 'video'
    output_composition = "output/composition." + variant + "." + file_name
    output_alpha = "output/alpha." + variant + "." + file_name
    output_foreground = "output/foreground." + variant + "." + file_name
    output_video_mbps = 4
    seq_chunk = 1

    # 初始化模型
    model = MattingNetwork(variant).eval().to(device)
    model.load_state_dict(torch.load(checkpoint, map_location=device))
    model = torch.jit.script(model)
    model = torch.jit.freeze(model)

    # 调用convert_video函数
    convert_video(
        model=model,
        input_source=input_source,
        output_type=output_type,
        output_composition=output_composition,
        output_alpha=output_alpha,
        output_foreground=output_foreground,
        output_video_mbps=output_video_mbps,
        seq_chunk=seq_chunk,
        device=device,
        dtype=torch.float32
    )


if __name__ == "__main__":
    main()
